"""
数据集下载服务
功能：处理RAVDESS等数据集的下载和组织
使用场景：模型训练前的数据准备
"""

import os
import logging
import shutil
from typing import Dict, Any
from pathlib import Path

logger = logging.getLogger(__name__)

class DatasetDownloader:
    """数据集下载服务类"""
    
    def __init__(self, datasets_dir: str = "datasets", config_file: str = "config/models_config.json"):
        """
        初始化下载服务
        参数：
            datasets_dir - 数据集存储目录
            config_file - 配置文件路径
        """
        self.datasets_dir = Path(datasets_dir)
        self.datasets_dir.mkdir(exist_ok=True)
        
        # 加载配置文件
        self.config = self._load_config(config_file)
        
        # 从配置文件读取RAVDESS数据集配置
        self.ravdess_config = self.config["ravdess_dataset"].copy()
        self.ravdess_config["local_path"] = self.datasets_dir / "ravdess"

    def _load_config(self, config_file: str) -> dict:
        """
        加载JSON配置文件
        参数：config_file - 配置文件路径
        返回：配置字典
        """
        import json
        
        config_path = Path(config_file)
        if not config_path.exists():
            config_path = Path(__file__).parent.parent / config_file
        
        if not config_path.exists():
            raise FileNotFoundError(f"配置文件不存在: {config_file}")
        
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except json.JSONDecodeError as e:
            raise ValueError(f"配置文件格式错误: {e}")
        except Exception as e:
            raise RuntimeError(f"加载配置文件失败: {e}")

    async def download_ravdess_auto(self) -> Dict[str, Any]:
        """
        自动下载RAVDESS数据集 - 仅使用Kaggle API
        返回：下载结果信息
        """
        logger.info("开始下载RAVDESS数据集...")
        
        # 检查是否已存在
        if self._check_ravdess_exists():
            return {
                "status": "success", 
                "message": "RAVDESS数据集已存在",
                "path": str(self.ravdess_config["local_path"]),
                "file_count": self._count_ravdess_files()
            }

        # 使用Kaggle API下载
        result = await self._try_kaggle_download()
        if result["status"] == "success":
            return result

        # 下载失败，返回手动下载指导
        kaggle_url = f"https://www.kaggle.com/datasets/{self.ravdess_config['kaggle_dataset']}"
        return {
            "status": "manual_required",
            "message": "自动下载失败，需要手动下载",
            "instructions": [
                f"请按以下步骤手动下载{self.ravdess_config['name']}数据集：",
                "",
                f"1. 访问 {kaggle_url}",
                "2. 点击 'Download' 下载数据集",
                "3. 解压下载的zip文件到本地",
                "4. 将所有.wav文件复制到 datasets/ravdess/ 目录下",
                "",
                "完成后重新点击训练按钮"
            ],
            "expected_path": str(self.ravdess_config["local_path"]),
            "dataset_info": {
                "name": self.ravdess_config["name"],
                "description": self.ravdess_config["description"],
                "total_samples": self.ravdess_config.get("total_samples", "未知"),
                "emotions": self.ravdess_config["all_emotions"]
            }
        }

    async def _try_kaggle_download(self) -> Dict[str, Any]:
        """尝试使用Kaggle API下载"""
        try:
            import kaggle
            
            # 检查Kaggle API配置
            kaggle.api.authenticate()
            
            logger.info("使用Kaggle API下载RAVDESS数据集...")
            
            # 下载数据集
            kaggle.api.dataset_download_files(
                self.ravdess_config["kaggle_dataset"],
                path=str(self.datasets_dir / "temp"),
                unzip=True
            )
            
            # 组织文件结构
            await self._organize_ravdess_files(self.datasets_dir / "temp")
            
            # 清理临时文件
            shutil.rmtree(self.datasets_dir / "temp", ignore_errors=True)
            
            return {
                "status": "success",
                "message": "RAVDESS数据集下载成功",
                "method": "kaggle_api",
                "path": str(self.ravdess_config["local_path"]),
                "file_count": self._count_ravdess_files()
            }
            
        except ImportError:
            logger.warning("Kaggle包未安装")
            return {"status": "failed", "reason": "kaggle_not_installed"}
        except Exception as e:
            logger.warning(f"Kaggle API下载失败: {e}")
            return {"status": "failed", "reason": f"kaggle_error: {e}"}

    async def _organize_ravdess_files(self, source_dir: Path):
        """
        组织RAVDESS文件结构
        参数：source_dir - 源文件目录
        """
        # 创建目标目录
        target_dir = self.ravdess_config["local_path"]
        target_dir.mkdir(parents=True, exist_ok=True)
        
        # 查找所有.wav文件
        wav_files = list(source_dir.rglob("*.wav"))
        
        logger.info(f"找到 {len(wav_files)} 个音频文件")
        
        # 复制文件到目标目录
        for wav_file in wav_files:
            target_file = target_dir / wav_file.name
            shutil.copy2(wav_file, target_file)
            logger.debug(f"复制文件: {wav_file.name}")

    def _check_ravdess_exists(self) -> bool:
        """
        检查RAVDESS数据集是否已存在
        返回：是否存在
        """
        ravdess_dir = self.ravdess_config["local_path"]
        if not ravdess_dir.exists():
            return False
        
        # 检查是否有.wav文件
        wav_files = list(ravdess_dir.glob("*.wav"))
        return len(wav_files) > 0

    def _count_ravdess_files(self) -> int:
        """
        统计RAVDESS数据集文件数量
        返回：文件数量
        """
        ravdess_dir = self.ravdess_config["local_path"]
        if not ravdess_dir.exists():
            return 0
        
        wav_files = list(ravdess_dir.glob("*.wav"))
        return len(wav_files)

    async def get_dataset_status(self) -> Dict[str, Any]:
        """
        获取数据集状态
        返回：状态信息
        """
        return {
            "ravdess": {
                "exists": self._check_ravdess_exists(),
                "path": str(self.ravdess_config["local_path"]),
                "file_count": self._count_ravdess_files(),
                "config": {
                    "name": self.ravdess_config["name"],
                    "description": self.ravdess_config["description"],
                    "kaggle_dataset": self.ravdess_config["kaggle_dataset"],
                    "emotions": self.ravdess_config["all_emotions"]
                }
            }
        } 